# Licensed under the MIT license.

"""
IO_System 模块

该模块提供了一个统一的输入/输出系统，用于与不同的语言模型API进行交互。
支持多种API类型包括vLLM、OpenAI GPT-3.5-turbo和调试模式。

主要功能：
- 统一的模型推理接口
- 支持批量和单个输入处理
- 自动计数API调用次数和token使用量
- 灵活的生成参数配置
"""

import sys

sys.path.append(".")

from typing import List, Dict

# 尝试导入vLLM API，如果导入失败则忽略
try:
    from models.vLLM_API import generate_with_vLLM_model
except:
    pass

# 尝试导入OpenAI API，如果导入失败则忽略
try:
    from models.OpenAI_API import generate_n_with_OpenAI_model
except:
    pass


class IO_System:
    """
    输入/输出系统类
    
    该类提供了一个统一的接口来与不同的语言模型API进行交互，
    支持vLLM、OpenAI GPT-3.5-turbo和调试模式。
    
    Attributes:
        api (str): 使用的API类型 ("vllm", "gpt3.5-turbo", "together", "debug")
        model_ckpt (str): 模型检查点路径或名称
        temperature (float): 生成温度参数，控制输出的随机性
        top_k (int): Top-K采样参数
        top_p (float): Top-P采样参数
        tokenizer: 分词器对象
        model: 模型对象
        call_counter (int): API调用计数器
        token_counter (int): Token使用计数器
    """

    def __init__(self, args, tokenizer, model) -> None:
        """
        初始化IO_System实例
        
        Args:
            args: 包含配置参数的对象，必须包含以下属性：
                - api: API类型
                - model_ckpt: 模型检查点
                - temperature: 生成温度
                - top_k: Top-K参数
                - top_p: Top-P参数
            tokenizer: 分词器对象（某些API类型下可为None）
            model: 模型对象（某些API类型下可为None或字符串）
        """
        self.api = args.api
        
        # 根据API类型验证参数
        if self.api == "together":
            assert tokenizer is None and model is None
        elif self.api == "gpt3.5-turbo":
            assert tokenizer is None and isinstance(model, str)
            
        # 存储模型配置参数
        self.model_ckpt = args.model_ckpt
        self.temperature = args.temperature
        self.top_k = args.top_k
        self.top_p = args.top_p
        self.tokenizer = tokenizer
        self.model = model

        # 初始化计数器
        self.call_counter = 0  # API调用次数计数
        self.token_counter = 0  # Token使用量计数

    def generate(self, model_input, max_tokens: int, num_return: int, stop_tokens):
        """
        生成文本的核心方法
        
        该方法根据配置的API类型调用相应的生成函数，支持单个字符串输入和批量列表输入。
        
        Args:
            model_input (str or List[str]): 输入文本，可以是单个字符串或字符串列表
            max_tokens (int): 生成的最大token数量
            num_return (int): 每个输入返回的生成结果数量
            stop_tokens: 停止生成的token列表
            
        Returns:
            List: 生成的文本结果列表
            - 如果输入是字符串，返回List[str]
            - 如果输入是列表，返回List[List[str]]
            
        Raises:
            NotImplementedError: 当API类型不支持时抛出
        """
        # 处理单个字符串输入
        if isinstance(model_input, str):
            if self.api == "vllm":
                # 使用vLLM API生成
                vllm_response = generate_with_vLLM_model(
                    self.model,
                    input=model_input,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    top_k=self.top_k,
                    n=num_return,
                    max_tokens=max_tokens,
                    stop=stop_tokens,
                )
                # 提取生成的文本
                io_output_list = [o.text for o in vllm_response[0].outputs]
                # 更新计数器
                self.call_counter += 1
                self.token_counter += sum([len(o.token_ids) for o in vllm_response[0].outputs])
                
            elif self.api == "gpt3.5-turbo":
                # 使用OpenAI GPT-3.5-turbo API生成
                gpt_response = generate_n_with_OpenAI_model(
                    prompt=model_input,
                    n=num_return,
                    model_ckpt=self.model,
                    max_tokens=max_tokens,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    top_k=self.top_k,
                    stop=["\n", "Answer"],  # 固定的停止token
                )
                io_output_list = gpt_response
                # 更新计数器（OpenAI API的token计数暂时设为0）
                self.call_counter += num_return
                self.token_counter += 0
                
            elif self.api == "debug":
                # 调试模式，返回固定的测试文本
                io_output_list = ["Debug: The answer is generated with debug mode, 233." for _ in range(num_return)]
            else:
                raise NotImplementedError(f"API {self.api} is not implemented.")
                
        # 处理批量列表输入
        elif isinstance(model_input, list):
            if self.api == "vllm":
                # 使用vLLM API批量生成
                vllm_response = generate_with_vLLM_model(
                    self.model,
                    input=model_input,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    top_k=self.top_k,
                    n=num_return,
                    max_tokens=max_tokens,
                    stop=stop_tokens,
                )
                # 为每个输入提取对应的生成结果
                io_output_list = [
                    [o.text for o in resp_to_single_input.outputs] for resp_to_single_input in vllm_response
                ]
                # 更新计数器
                self.call_counter += 1
                self.token_counter += sum(
                    [
                        sum([len(o.token_ids) for o in resp_to_single_input.outputs])
                        for resp_to_single_input in vllm_response
                    ]
                )
                
            elif self.api == "gpt3.5-turbo":
                # 使用OpenAI API逐个处理输入
                io_output_list = []
                for input in model_input:
                    gpt_response = generate_n_with_OpenAI_model(
                        prompt=input,
                        n=num_return,
                        model_ckpt=self.model,
                        max_tokens=max_tokens,
                        temperature=self.temperature,
                        top_p=self.top_p,
                        top_k=self.top_k,
                        stop=["\n", "Answer"],  # 固定的停止token
                    )
                    io_output_list.append(gpt_response)
                    # 为每个输入更新计数器
                    self.call_counter += num_return
                    self.token_counter += 0
                    
            elif self.api == "debug":
                # 调试模式，为每个输入返回固定的测试文本
                io_output_list = [
                    ["Debug: The answer is generated with debug mode, 233." for _ in range(num_return)]
                    for _ in model_input
                ]
            else:
                raise NotImplementedError(f"API {self.api} is not implemented.")

        return io_output_list
